#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   socketServer.py
@Time    :   2024-03-25 23:40:56
@Author  :   chakcy 
@Email   :   947105045@qq.com
@description   :   SocketServer类
'''

import socket  
import threading  
from common.B64 import base64_decode, base64_encode  

# 多线程SocketServer类  
class SocketServer:  
    def __init__(self, ip, port):  
        self.ip = ip  
        self.port = port  
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)  
        self.server_socket.bind((self.ip, self.port))  
        self.server_socket.listen(5)  
        self.clients = {}  # 用于存储客户端连接的字典  
        self.lock = threading.Lock()
  
    def start(self, process_data):  
        try:  
            print(f"Server started on {self.ip}:{self.port}")  
            while True:  
                client_socket, client_addr = self.server_socket.accept()  
                print(f"Accepted connection from {client_addr}")  
                self.lock.acquire()  
                self.clients[client_socket] = client_addr
                self.lock.release()
                threading.Thread(target=self.handle_client, args=(client_socket, process_data)).start()  
                # self.handle_client(client_socket, process_data)  
        except KeyboardInterrupt:  
            print("Server stopped")  
        finally:  
            self.server_socket.close()  
            for client_socket in self.clients:  
                client_socket.close()  
  
    def handle_client(self, client_socket, process_data):  
        client_addr = self.clients[client_socket]  # 获取客户端地址，假设此时client_socket还在clients字典中  
        try:  
            while True:  
                data = client_socket.recv(1024)  
                if not data:  
                    break  
                decoded_data = base64_decode(data)  
                print(f"Received from {client_addr}: {decoded_data}")  
                # 假设有一个处理函数叫 process_data  
                result = base64_encode(process_data(decoded_data))
                client_socket.sendall(result.encode('utf-8'))  
        except Exception as e:  
            print(f"Error handling client {client_addr}: {e}")  
        finally:  
            # 在关闭socket之前，从clients字典中移除对应的项  
            self.lock.acquire()  
            if client_socket in self.clients:  
                del self.clients[client_socket]
            self.lock.release()  
            client_socket.close()  # 关闭socket  
            print(f"Connection closed with {client_addr}") 

    def close(self):
        print("Closing all client connections...")
        for client_socket in list(self.clients.keys()):
            client_socket.close()
            del self.clients[client_socket]
        print("All client connections closed.")
        print("Closing server connection...")
        self.server_socket.close()
        print("Server connection closed.")

import asyncio  
import functools

# 异步SocketServer类  
class AsyncSocketServer:
    def __init__(self, ip, port):
        self.ip = ip
        self.port = port
        self.server = None

    async def start(self, process_data):
        self.server = await asyncio.start_server(
            functools.partial(self.handle_client, process_data=process_data),
            self.ip, 
            self.port,
        )
        print(f"Server started on {self.ip}:{self.port}")

        async with self.server:
            await self.server.serve_forever()
    
    async def handle_client(self, reader, writer, process_data):  
        try:  
            while True:  
                data = await reader.read(1024)  
                if not data:  
                    break  
                decoded_data = base64_decode(data)  
                print(f"Received data: {decoded_data}")  
                result = base64_encode(process_data(decoded_data))  
                writer.write(result.encode('utf-8'))  
                await writer.drain()  
        except Exception as e:  
            print(f"Error handling client: {e}")  
        finally:  
            writer.close()  
            await writer.wait_closed()
  
    async def close(self):  
        if self.server is not None:  
            self.server.close()  
            await self.server.wait_closed()
            
# 使用示例  
if __name__ == "__main__":  
    def process_data(data):  
        return data + " processed"  
    # server = SocketServer('127.0.0.1', 8888)  
    # server.start(process_data)

    async def main():
        server = AsyncSocketServer('127.0.0.1', 8888)
        await server.start(process_data)
        await server.close()

    asyncio.run(main())
